import backtrader as bt
import pandas as pd
import numpy as np
import logging
import yaml
import os
from datetime import datetime
from .data_loader import CustomCSVData


class BacktestEngine:
    """
    回测引擎类，负责加载配置、数据和执行回测
    """

    def __init__(self, config_file='config.yaml'):
        """
        初始化回测引擎

        参数:
            config_file: 配置文件路径
        """
        self.config_file = config_file
        self.cerebro = bt.Cerebro()
        self.logger = logging.getLogger(__name__)
        self.config = None
        self.load_config()

        # 设置Cerebro参数
        self.cerebro.broker.setcash(self.config.get('initial_cash', 1000000))
        self.cerebro.broker.setcommission(
            commission=self.config.get('commission', 0.0005)
        )

        # 添加分析器
        self._add_analyzers()

    def load_config(self):
        """
        加载配置文件
        """
        try:
            self.logger.info(f"加载配置文件: {self.config_file}")
            with open(self.config_file, 'r') as f:
                self.config = yaml.safe_load(f)
            self.logger.info("配置文件加载成功")
        except Exception as e:
            self.logger.error(f"加载配置文件失败: {str(e)}")
            raise

    def _add_analyzers(self):
        """
        添加分析器到Cerebro
        """
        self.cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='sharpe_ratio')
        self.cerebro.addanalyzer(bt.analyzers.DrawDown, _name='drawdown')
        self.cerebro.addanalyzer(bt.analyzers.Returns, _name='returns')
        self.cerebro.addanalyzer(bt.analyzers.TradeAnalyzer, _name='trades')
        self.cerebro.addanalyzer(bt.analyzers.Transactions, _name='transactions')
        self.logger.info("已添加分析器")

    def load_data(self, file_path):
        """
        加载单个资产数据

        参数:
            file_path: 数据文件路径

        返回:
            tuple: (数据DataFrame, 开始日期, 结束日期)
        """
        try:
            self.logger.info(f"加载数据文件: {file_path}")
            dtformat = self.config.get('dtformat', '%Y/%m/%d')
            df = CustomCSVData.load_data(file_path, dtformat)

            # 检查数据是否为空
            if df.empty:
                self.logger.error(f"加载的数据为空: {file_path}")
                raise ValueError(f"数据文件为空: {file_path}")

            start_date = df.index.min()
            end_date = df.index.max()
            self.logger.info(f"数据加载成功: {start_date} 至 {end_date}, 共 {len(df)} 条记录")
            return df, start_date, end_date
        except Exception as e:
            self.logger.error(f"加载数据文件失败: {str(e)}")
            raise

    def add_data_to_cerebro(self, df, name):
        """
        将数据添加到Cerebro

        参数:
            df: 数据DataFrame
            name: 数据名称
        """
        try:
            # 创建PandasData对象
            data = bt.feeds.PandasData(
                dataname=df,
                datetime=None,  # 使用索引作为日期
                open=1,
                high=2,
                low=3,
                close=0,
                volume=4,
                openinterest=-1
            )

            # 设置数据名称
            data._name = name

            # 添加到Cerebro
            self.cerebro.adddata(data)
            self.logger.info(f"已添加数据源: {name}")
        except Exception as e:
            self.logger.error(f"添加数据到Cerebro失败: {str(e)}")
            raise

    def run(self, strategy):
        """
        执行回测

        参数:
            strategy: 策略类

        返回:
            list: 回测结果
        """
        try:
            # 检查配置文件
            if 'assets' not in self.config:
                raise ValueError("配置文件中缺少'assets'部分")

            # 加载所有资产数据
            datas = []
            for asset in self.config['assets']:
                if 'id' not in asset or 'file' not in asset:
                    raise ValueError("资产配置中缺少'id'或'file'字段")

                file_path = asset['file']
                if not os.path.exists(file_path):
                    raise FileNotFoundError(f"文件不存在: {file_path}")

                df, start, end = self.load_data(file_path)
                datas.append({
                    'df': df,
                    'start': start,
                    'end': end,
                    'name': asset['id']
                })

            # 获取所有数据源的共同日期范围
            all_starts = [data['start'] for data in datas]
            all_ends = [data['end'] for data in datas]
            common_start = max(all_starts)
            common_end = min(all_ends)

            # 检查共同日期范围是否有效
            if common_start > common_end:
                self.logger.error(f"无效的共同日期范围: 开始日期 {common_start} 晚于结束日期 {common_end}")
                self.logger.error("可能原因: 资产数据没有重叠的日期范围")
                return None

            self.logger.info(f"共同日期范围: {common_start} 至 {common_end}")

            # 截取共同日期范围并添加到Cerebro
            valid_data_count = 0
            for data in datas:
                aligned_df = data['df'].loc[common_start:common_end]

                # 检查截取后的数据是否为空
                if aligned_df.empty:
                    self.logger.error(f"资产 {data['name']} 在共同日期范围内无数据: {common_start} 至 {common_end}")
                    continue

                self.add_data_to_cerebro(aligned_df, data['name'])
                valid_data_count += 1

            # 检查是否有有效数据
            if valid_data_count == 0:
                self.logger.error("所有资产在共同日期范围内均无有效数据，回测终止")
                return None

            # 添加策略
            self.cerebro.addstrategy(strategy)
            self.logger.info(f"已添加策略: {strategy.__name__}")

            # 运行回测
            self.logger.info("开始回测...")
            results = self.cerebro.run()
            self.logger.info("回测完成")

            return results
        except Exception as e:
            self.logger.error(f"回测失败: {str(e)}")

            # 添加特定错误类型的诊断信息
            if "division by zero" in str(e):
                self.logger.error("可能原因: 回测期间无有效数据或时间周期为零")
                self.logger.error("解决方案: 检查资产数据的日期范围和重叠区域")
            elif "out of bounds" in str(e):
                self.logger.error("可能原因: 数据索引或日期范围错误")
                self.logger.error("解决方案: 检查数据加载和日期对齐逻辑")

            # 打印详细错误信息
            import traceback
            traceback.print_exc()
            return None

    def get_results_summary(self, result):
        """
        获取回测结果摘要

        参数:
            result: 回测结果对象

        返回:
            dict: 结果摘要
        """
        if not result:
            return {}

        # 安全获取分析器结果，避免分析器本身出错
        try:
            sharpe = result.analyzers.sharpe_ratio.get_analysis() if hasattr(result.analyzers, 'sharpe_ratio') else {}
            drawdown = result.analyzers.drawdown.get_analysis() if hasattr(result.analyzers, 'drawdown') else {}
            returns = result.analyzers.returns.get_analysis() if hasattr(result.analyzers, 'returns') else {}
            trades = result.analyzers.trades.get_analysis() if hasattr(result.analyzers, 'trades') else {}
            transactions = result.analyzers.transactions.get_analysis() if hasattr(result.analyzers,
                                                                                   'transactions') else {}
        except Exception as e:
            self.logger.error(f"获取分析器结果失败: {str(e)}")
            sharpe = {}
            drawdown = {}
            returns = {}
            trades = {}
            transactions = {}

        summary = {
            'final_value': self.cerebro.broker.getvalue(),
            'sharpe_ratio': sharpe,
            'drawdown': drawdown,
            'returns': returns,
            'trades': trades,
            'transactions': transactions
        }

        return summary

    def plot_results(self):
        """
        绘制回测结果图表
        """
        try:
            self.logger.info("绘制回测结果图表...")
            self.cerebro.plot()
            self.logger.info("图表绘制完成")
        except Exception as e:
            self.logger.error(f"绘制图表失败: {str(e)}")